import numpy as np
from matplotlib.pyplot import rcParams, figure, imshow, title, xlabel, ylabel, colorbar, contourf, show
from PIL.Image import open, ANTIALIAS, FLIP_LEFT_RIGHT, FLIP_TOP_BOTTOM, ROTATE_90, ROTATE_270
from numpy import array, min, max, arctan2, arctan, pi, power, NaN, zeros, tan

rcParams['font.sans-serif'] = ['SimHei']
rcParams['axes.unicode_minus'] = False

class Redraw_AoP:
    def __init__(self, img_path1, img_path2, img_path3, flip_image):
        self.img1 = img_path1
        self.img2 = img_path2
        self.img3 = img_path3
        self.flip_img = flip_image

        NAME_FORWARD = self.img1[-26: -7]
        self.SAVE_NAME = "./REASULT_POLARIZATION_DISTRIBUTION_IMAGE/" + NAME_FORWARD
    def Redraw_AoP1(self):

        img1 = open(self.img1)
        img2 = open(self.img2)
        img3 = open(self.img3)

        flag = img1.mode == 'I;16'

        if self.flip_img == 'xEyN':
            img1 = img1.transpose(FLIP_TOP_BOTTOM)
            img2 = img2.transpose(FLIP_TOP_BOTTOM)
            img3 = img3.transpose(FLIP_TOP_BOTTOM)

        elif self.flip_img == 'xNyW':
            img1 = img1.transpose(ROTATE_270)
            img2 = img2.transpose(ROTATE_270)
            img3 = img3.transpose(ROTATE_270)
            img1 = img1.transpose(FLIP_LEFT_RIGHT)
            img2 = img2.transpose(FLIP_LEFT_RIGHT)
            img3 = img3.transpose(FLIP_LEFT_RIGHT)

        elif self.flip_img == 'xSyE':
            img11 = img1.transpose(ROTATE_90)
            img22 = img2.transpose(ROTATE_90)
            img33 = img3.transpose(ROTATE_90)
            img1 = img11.transpose(FLIP_LEFT_RIGHT)
            img2 = img22.transpose(FLIP_LEFT_RIGHT)
            img3 = img33.transpose(FLIP_LEFT_RIGHT)

        elif self.flip_img == 'xWyS':
            img1 = img1.transpose(FLIP_LEFT_RIGHT)
            img2 = img2.transpose(FLIP_LEFT_RIGHT)
            img3 = img3.transpose(FLIP_LEFT_RIGHT)

        else:
            img1 = img1
            img2 = img2
            img3 = img3

        if flag:
            img1 = array(img1)
            img2 = array(img2)
            img3 = array(img3)
        else:
            img1 = array(img1.convert('L'))
            img2 = array(img2.convert('L'))
            img3 = array(img3.convert('L'))

        width = img1.shape[0]
        high = img1.shape[1]
        zoom_out = 0.25
        # if width > 640 :
        #     img1 = np.resize(img1, (width * zoom_out, high * zoom_out))
        #     img2 = np.resize(img2, (width * zoom_out, high * zoom_out))
        #     img3 = np.resize(img3, (width * zoom_out, high * zoom_out))
        # else :
        #     img1 = img1
        #     img2 = img2
        #     img3 = img3

        if not flag:
            src1_im2double = img1 / 255
            src120_mat2gray = (src1_im2double - min(src1_im2double)) / (
                    max(src1_im2double) - min(src1_im2double))
            src2_im2double = img2 / 255
            src0_mat2gray = (src2_im2double - min(src2_im2double)) / (
                    max(src2_im2double) - min(src2_im2double))
            src3_im2double = img3 / 255
            src60_mat2gray = (src3_im2double - min(src3_im2double)) / (
                    max(src3_im2double) - min(src3_im2double))
        else:
            src0_mat2gray = img2
            src60_mat2gray = img3
            src120_mat2gray = img1

        I = (2 / 3) * (src0_mat2gray + src60_mat2gray + src120_mat2gray)
        Q = (2 / 3) * (2 * src0_mat2gray - src60_mat2gray - src120_mat2gray)
        U = ((2 * np.sqrt(3)) / 3) * (src60_mat2gray - src120_mat2gray)

        aop = 0.5 * arctan2(U, Q)
        dop = np.sqrt(np.power(Q, 2) + np.power(U, 2)) / I

        h = aop.shape[0]
        w = aop.shape[1]
        if h > w:
            r = w / 2
        else:
            r = h / 2
        aop_circle = aop
        dop_circle = dop

        for i in range(h):
            for j in range(w):
                if power((i - (h / 2)), 2) + power((j - (w / 2)), 2) > power(r, 2):
                    aop_circle[i, j] = NaN
                    dop_circle[i, j] = NaN

        azim_angle = zeros((h, w))
        aoe = zeros((h, w))
        for i in range(0, h):
            for j in range(0, w):
                if power((i - h / 2), 2) + power((j - w / 2), 2) <= power(r, 2):
                    if j == w / 2:
                        azim_angle[i, j] = pi / 2
                    else:
                        azim_angle[i, j] = arctan((i - h / 2) / (j - w / 2))

                    aoe[i, j] = arctan(tan(aop[i, j] - azim_angle[i, j])) * 180 / pi
                    if -90 <= aoe[i, j] <= 90:
                        aoe[i, j] = aoe[i, j] - 90
                    elif aoe[i, j] > 90:
                        aoe[i, j] = aoe[i, j] - 180
                    elif -90 < aoe[i, j]:
                        aoe[i, j] = aoe[i, j] + 180
                else:
                    aoe[i, j] = NaN

        return aoe, aop, dop


if __name__ == '__main__':
    path1 = r'E:\git_code\水下模型建立\img\1-0.tiff'
    path2 = r'E:\git_code\水下模型建立\img\1-73.78057111627271.tiff'
    path3 = r'E:\git_code\水下模型建立\img\1-147.56114223254545.tiff'
    draw_img = Redraw_AoP(path3, path1, path2, None)
    aoe, aop, dop = draw_img.Redraw_AoP1()


